    

from sfepy.discrete import Region
from sfepy.discrete.fem import FEDomain

import numpy as np



def get_cavity_region(domain: FEDomain,
                      omega: Region, 
                      target_corr,
                      side_face: Region,
                      axis: int):
    """
    根据指定的坐标位置沿着指定的方向从指定点到指定边缘去除线段上的单元。本函数目前仅适用于8节点6面体单元。

    参数
    ----------
    domain : FEDomain
        需要分析的几何结构的FEDomain对象。
    omega : Region
        表征domain全局区域的Region，通常通过domain.create_region('Omega', 'all')获取。
    target_corr : ndarray
        需要去除的目标单元的坐标，ndarray.shape=(1,3)。
    side_face : Region
        从target_corr开始，沿着axis所只是的方向将路径上经过的单元全部去除。
    axis : int
        1表示x轴正方向，2表示y轴正方向，3表示z轴正方向。
        -1表示x轴负方向，-2表示y轴负方向，-3表示z轴负方向。

    返回值
    ---------
    cavity_cellid : 需要去除的单元的cellid
    """
    cells_nodes = domain.get_conn()
    node_corr = omega.cmesh.coors
    cells_nodes_corr = node_corr[cells_nodes]
    cells_center = np.mean(cells_nodes_corr, axis=1)

    cavity_cellid = np.argmin(np.linalg.norm(cells_center - target_corr, axis=1))
    cavity_cellid = np.array([cavity_cellid])

    
    while (1):
        cell_node_corr = node_corr[cells_nodes[cavity_cellid[-1]]]
        if axis > 0:
            side_node_indices = cells_nodes[cavity_cellid[-1], 
                                            np.argpartition(cell_node_corr[:, abs(axis) - 1], -4)[-4:]]
        else:
            side_node_indices = cells_nodes[cavity_cellid[-1], 
                                            np.argpartition(cell_node_corr[:, abs(axis) - 1], 4)[:4]]
        
        # 判断cavity cell是否已经到达了边缘，如果是则结束遍历并输出结果
        if np.isin(side_node_indices, side_face.vertices).all():
            break

        cell_mask = np.ones(cells_nodes.shape[0], dtype=bool)
        cell_mask[cavity_cellid] = False   # 设置单元掩码，避免遍历重复的单元
        getnext_tag = False
        for cell_indice, cell_nodes in enumerate(cells_nodes):
            if np.isin(side_node_indices, cell_nodes).all() and not np.isin(cell_indice, cavity_cellid):
                cavity_cellid = np.append(cavity_cellid, [cell_indice])
                getnext_tag = True
        if getnext_tag == False:
            raise ValueError('Cavity cells not found !!!')
    
    return cavity_cellid
    






